package cn.bitleo.app.src.main.java.Injury;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtSession;
import ai.onnxruntime.OrtException;

import java.nio.FloatBuffer;
import java.util.Collections;

public class Injury3DModule{

    private final OrtEnvironment env;
    private final OrtSession session;

    public Injury3DModule(OrtEnvironment env, OrtSession session) {
        this.env = env;
        this.session = session;
    }

    /**
     * 输入单样本一维数组 (长度 = 5)
     * 输出 float[1][5] (batch=1)
     */
    public float[][] inferFrom1D(float[] features) throws OrtException {
        if (features.length != 5) {
            throw new IllegalArgumentException("Input size must be exactly 5");
        }

        // 转成 [1,1,5] 给 ONNX
        long[] shape = new long[]{5};
        OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(features), shape);

        // 模型输入名
        String inputName = session.getInputNames().iterator().next();

        // 推理
        OrtSession.Result result = session.run(Collections.singletonMap(inputName, tensor));

        // 获取输出
        Object value = result.get(0).getValue();

        // 输出 shape = (batch,5)
        if (value instanceof float[][]) {
            return (float[][]) value;  // shape = [1][5]
        } else if (value instanceof float[]) {
            // 处理异常情况，单样本返回一维数组
            float[] out = (float[]) value;
            return new float[][]{out};
        } else {
            throw new RuntimeException("Unexpected output type: " + value.getClass().getName());
        }
    }

    /**
     * 输入多样本 batch
     * 输入 shape = [batch, 5]
     * 输出 shape = [batch][5]
     */
    public float[][] inferBatch(float[][] batchFeatures) throws OrtException {
        int batchSize = batchFeatures.length;

        // flatten 为 [batch*5]
        float[] flat = new float[batchSize * 5];
        for (int i = 0; i < batchSize; i++) {
            if (batchFeatures[i].length != 5) {
                throw new IllegalArgumentException("Each sample must have length 5");
            }
            System.arraycopy(batchFeatures[i], 0, flat, i * 5, 5);
        }

        // reshape [batch,1,5]
        long[] shape = new long[]{batchSize, 1, 5};
        OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(flat), shape);

        String inputName = session.getInputNames().iterator().next();
        OrtSession.Result result = session.run(Collections.singletonMap(inputName, tensor));

        Object value = result.get(0).getValue();
        if (value instanceof float[][]) {
            return (float[][]) value; // shape = [batch][5]
        } else {
            throw new RuntimeException("Unexpected output type: " + value.getClass().getName());
        }
    }
}
